#!/usr/bin/env python3
"""
验证灰度图训练模式是否正确配置

运行方式：
    python test_grayscale_mode.py
"""

import torch
import numpy as np
from model import DexiNed

def test_model_channels():
    """测试模型在不同输入通道下的工作状态"""
    print("=" * 60)
    print("测试 DexiNed 模型通道配置")
    print("=" * 60)
    
    # 测试灰度图模型
    print("\n1. 测试灰度图模型 (in_channels=1):")
    try:
        model_gray = DexiNed(in_channels=1)
        x_gray = torch.randn(2, 1, 640, 640)
        
        model_gray.eval()
        with torch.no_grad():
            outputs_gray = model_gray(x_gray)
        
        print(f"   ✓ 输入形状: {x_gray.shape}")
        print(f"   ✓ 输出数量: {len(outputs_gray)} 个特征图")
        print(f"   ✓ 第一个输出形状: {outputs_gray[0].shape}")
        print(f"   ✓ 最后一个输出形状: {outputs_gray[-1].shape}")
        
        # 统计参数量
        total_params = sum(p.numel() for p in model_gray.parameters())
        trainable_params = sum(p.numel() for p in model_gray.parameters() if p.requires_grad)
        print(f"   ✓ 总参数量: {total_params:,}")
        print(f"   ✓ 可训练参数: {trainable_params:,}")
        
        # 检查第一层卷积
        first_conv_weight = model_gray.block_1.conv1.weight
        print(f"   ✓ 第一层卷积权重形状: {first_conv_weight.shape}")
        assert first_conv_weight.shape[1] == 1, "第一层输入通道应为1"
        
        print("   ✓ 灰度图模型测试通过！")
    except Exception as e:
        print(f"   ✗ 灰度图模型测试失败: {e}")
        return False
    
    # 测试彩色图模型
    print("\n2. 测试彩色图模型 (in_channels=3):")
    try:
        model_color = DexiNed(in_channels=3)
        x_color = torch.randn(2, 3, 640, 640)
        
        model_color.eval()
        with torch.no_grad():
            outputs_color = model_color(x_color)
        
        print(f"   ✓ 输入形状: {x_color.shape}")
        print(f"   ✓ 输出数量: {len(outputs_color)} 个特征图")
        print(f"   ✓ 第一个输出形状: {outputs_color[0].shape}")
        print(f"   ✓ 最后一个输出形状: {outputs_color[-1].shape}")
        
        # 统计参数量
        total_params = sum(p.numel() for p in model_color.parameters())
        trainable_params = sum(p.numel() for p in model_color.parameters() if p.requires_grad)
        print(f"   ✓ 总参数量: {total_params:,}")
        print(f"   ✓ 可训练参数: {trainable_params:,}")
        
        # 检查第一层卷积
        first_conv_weight = model_color.block_1.conv1.weight
        print(f"   ✓ 第一层卷积权重形状: {first_conv_weight.shape}")
        assert first_conv_weight.shape[1] == 3, "第一层输入通道应为3"
        
        print("   ✓ 彩色图模型测试通过！")
    except Exception as e:
        print(f"   ✗ 彩色图模型测试失败: {e}")
        return False
    
    # 比较参数量差异
    print("\n3. 参数量对比:")
    gray_params = sum(p.numel() for p in model_gray.parameters())
    color_params = sum(p.numel() for p in model_color.parameters())
    diff = color_params - gray_params
    diff_percent = (diff / color_params) * 100
    
    print(f"   灰度图模型: {gray_params:,} 参数")
    print(f"   彩色图模型: {color_params:,} 参数")
    print(f"   差异: {diff:,} 参数 ({diff_percent:.2f}%)")
    
    return True

def test_data_processing():
    """测试数据预处理流程"""
    print("\n" + "=" * 60)
    print("测试数据预处理")
    print("=" * 60)
    
    # 模拟灰度图数据处理
    print("\n1. 测试灰度图预处理:")
    try:
        # 模拟灰度图像
        gray_img = np.random.randint(0, 256, (640, 640), dtype=np.uint8)
        gray_img_3d = np.expand_dims(gray_img, axis=-1)  # HxWx1
        
        print(f"   ✓ 原始灰度图: {gray_img.shape}")
        print(f"   ✓ 扩展后: {gray_img_3d.shape}")
        
        # 模拟均值减法
        mean_bgr = [103.939, 116.779, 123.68]
        gray_mean = np.mean(mean_bgr)
        gray_img_normalized = gray_img_3d.astype(np.float32) - gray_mean
        
        print(f"   ✓ BGR均值: {mean_bgr}")
        print(f"   ✓ 灰度均值: {gray_mean:.3f}")
        print(f"   ✓ 归一化后范围: [{gray_img_normalized.min():.2f}, {gray_img_normalized.max():.2f}]")
        
        # 转换为CHW格式
        gray_img_chw = gray_img_normalized.transpose((2, 0, 1))
        print(f"   ✓ 转换为CHW: {gray_img_chw.shape}")
        
        print("   ✓ 灰度图预处理测试通过！")
    except Exception as e:
        print(f"   ✗ 灰度图预处理测试失败: {e}")
        return False
    
    # 模拟彩色图数据处理
    print("\n2. 测试彩色图预处理:")
    try:
        # 模拟彩色图像
        color_img = np.random.randint(0, 256, (640, 640, 3), dtype=np.uint8)
        
        print(f"   ✓ 原始彩色图: {color_img.shape}")
        
        # 模拟均值减法
        color_img_normalized = color_img.astype(np.float32) - np.array(mean_bgr)
        
        print(f"   ✓ BGR均值: {mean_bgr}")
        print(f"   ✓ 归一化后范围: [{color_img_normalized.min():.2f}, {color_img_normalized.max():.2f}]")
        
        # 转换为CHW格式
        color_img_chw = color_img_normalized.transpose((2, 0, 1))
        print(f"   ✓ 转换为CHW: {color_img_chw.shape}")
        
        print("   ✓ 彩色图预处理测试通过！")
    except Exception as e:
        print(f"   ✗ 彩色图预处理测试失败: {e}")
        return False
    
    return True

def test_forward_backward():
    """测试前向传播和反向传播"""
    print("\n" + "=" * 60)
    print("测试前向和反向传播")
    print("=" * 60)
    
    print("\n1. 测试灰度图前向反向传播:")
    try:
        model = DexiNed(in_channels=1)
        model.train()
        
        # 前向传播
        x = torch.randn(2, 1, 320, 320, requires_grad=True)
        outputs = model(x)
        
        print(f"   ✓ 输入: {x.shape}")
        print(f"   ✓ 输出数量: {len(outputs)}")
        
        # 模拟损失计算和反向传播
        loss = sum(out.sum() for out in outputs)
        loss.backward()
        
        print(f"   ✓ 损失值: {loss.item():.4f}")
        print(f"   ✓ 输入梯度: {x.grad is not None}")
        
        # 检查模型参数梯度
        has_grad = sum(1 for p in model.parameters() if p.grad is not None)
        total_params = sum(1 for _ in model.parameters())
        print(f"   ✓ 有梯度的参数: {has_grad}/{total_params}")
        
        print("   ✓ 前向反向传播测试通过！")
    except Exception as e:
        print(f"   ✗ 前向反向传播测试失败: {e}")
        return False
    
    return True

def main():
    """主测试函数"""
    print("\n" + "=" * 60)
    print("DexiNed 灰度图训练模式验证")
    print("=" * 60)
    
    results = []
    
    # 运行所有测试
    results.append(("模型通道配置", test_model_channels()))
    results.append(("数据预处理", test_data_processing()))
    results.append(("前向反向传播", test_forward_backward()))
    
    # 汇总结果
    print("\n" + "=" * 60)
    print("测试结果汇总")
    print("=" * 60)
    
    for test_name, passed in results:
        status = "✓ 通过" if passed else "✗ 失败"
        print(f"{test_name:20s}: {status}")
    
    all_passed = all(result[1] for result in results)
    
    print("\n" + "=" * 60)
    if all_passed:
        print("✓ 所有测试通过！灰度图训练模式已正确配置。")
        print("\n可以使用以下命令开始训练：")
        print("  python main.py --grayscale --train_data BSDS --img_width 640 --img_height 640")
    else:
        print("✗ 部分测试失败，请检查配置。")
    print("=" * 60 + "\n")
    
    return all_passed

if __name__ == '__main__':
    import sys
    success = main()
    sys.exit(0 if success else 1)
